{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = os.getenv('OPENAI_KEY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.graph_stores.neo4j import Neo4jGraphStore\n",
    "from llama_index.core import StorageContext\n",
    "\n",
    "username = \"neo4j\"\n",
    "password = os.getenv('NEO4J_PW')\n",
    "url = os.getenv('NEO4J_URL')\n",
    "database = \"neo4j\"\n",
    "\n",
    "graph_store = Neo4jGraphStore(\n",
    "    username=username,\n",
    "    password=password,\n",
    "    url=url,\n",
    "    database=database,\n",
    ")\n",
    "\n",
    "storage_context = StorageContext.from_defaults(graph_store=graph_store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.indices.knowledge_graph.retrievers import KnowledgeGraphRAGRetriever\n",
    "from llama_index.core.query_engine import RetrieverQueryEngine\n",
    "\n",
    "graph_rag_retriever_with_nl2graphquery = KnowledgeGraphRAGRetriever(\n",
    "    storage_context=storage_context,\n",
    "    verbose=True,\n",
    "    retriever_mode=\"keyword\",\n",
    "    with_nl2graphquery=True,\n",
    ")\n",
    "\n",
    "query_engine_with_nl2graphquery = RetrieverQueryEngine.from_args(\n",
    "    graph_rag_retriever_with_nl2graphquery,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<b>Harry Potter is a character in a popular series of fantasy novels written by J.K. Rowling.</b>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import display, Markdown\n",
    "\n",
    "response = query_engine_with_nl2graphquery.query(\n",
    "    \"who is harry potter\",\n",
    ")\n",
    "\n",
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'af4510aa-44a4-4720-aa83-838ef29eb60b': {'kg_rel_map': {'Harry': [['TURNED_OUT_TO_BE',\n",
      "                                                                    'Bag of '\n",
      "                                                                    'chips'],\n",
      "                                                                   ['ANGRILY',\n",
      "                                                                    'Said'],\n",
      "                                                                   ['HEART_TWANGING_LIKE',\n",
      "                                                                    'Giant '\n",
      "                                                                    'elastic '\n",
      "                                                                    'band'],\n",
      "                                                                   ['RAN',\n",
      "                                                                    'Before'],\n",
      "                                                                   ['LAY_AND_WATCHED',\n",
      "                                                                    'Birthday '\n",
      "                                                                    'tick '\n",
      "                                                                    'nearer'],\n",
      "                                                                   ['LAY_AND_WATCHED',\n",
      "                                                                    'Birthday'],\n",
      "                                                                   ['SHUFFLED_OFF_INTO',\n",
      "                                                                    'Kitchen'],\n",
      "                                                                   ['MOVED_OUT_OF',\n",
      "                                                                    'Cupboard'],\n",
      "                                                                   ['MOVED_INTO',\n",
      "                                                                    \"Dudley's \"\n",
      "                                                                    'second '\n",
      "                                                                    'bedroom'],\n",
      "                                                                   ['COULD_SEE',\n",
      "                                                                    'Uncle '\n",
      "                                                                    \"vernon's \"\n",
      "                                                                    'shiny '\n",
      "                                                                    'black '\n",
      "                                                                    'shoes'],\n",
      "                                                                   ['RECEIVED',\n",
      "                                                                    'First '\n",
      "                                                                    'letter'],\n",
      "                                                                   ['RECEIVED',\n",
      "                                                                    'Letter'],\n",
      "                                                                   ['PICKED_UP',\n",
      "                                                                    'Letter'],\n",
      "                                                                   ['STARED_AT',\n",
      "                                                                    'Letter'],\n",
      "                                                                   ['EARNED',\n",
      "                                                                    'Longest-ever '\n",
      "                                                                    'punishment'],\n",
      "                                                                   ['DREAMED_OF',\n",
      "                                                                    'Unknown '\n",
      "                                                                    'relation'],\n",
      "                                                                   ['FELL_HARD_ON',\n",
      "                                                                    'Concrete '\n",
      "                                                                    'floor'],\n",
      "                                                                   ['IMAGINED',\n",
      "                                                                    'School'],\n",
      "                                                                   ['SITTING_IN',\n",
      "                                                                    'Car'],\n",
      "                                                                   ['HATED',\n",
      "                                                                    'Mrs. '\n",
      "                                                                    'figg'],\n",
      "                                                                   ['TURNING_OVER',\n",
      "                                                                    'Bacon'],\n",
      "                                                                   ['WORE',\n",
      "                                                                    'Glasses'],\n",
      "                                                                   ['SLEPT_IN',\n",
      "                                                                    'Cupboard'],\n",
      "                                                                   ['HAD',\n",
      "                                                                    'Dream'],\n",
      "                                                                   ['HAD',\n",
      "                                                                    'Sleepless '\n",
      "                                                                    'night'],\n",
      "                                                                   ['HAD',\n",
      "                                                                    'Scar']]},\n",
      "                                          'kg_rel_text': [\"['TURNED_OUT_TO_BE', \"\n",
      "                                                          \"'Bag of chips']\",\n",
      "                                                          \"['ANGRILY', 'Said']\",\n",
      "                                                          \"['HEART_TWANGING_LIKE', \"\n",
      "                                                          \"'Giant elastic \"\n",
      "                                                          \"band']\",\n",
      "                                                          \"['RAN', 'Before']\",\n",
      "                                                          \"['LAY_AND_WATCHED', \"\n",
      "                                                          \"'Birthday tick \"\n",
      "                                                          \"nearer']\",\n",
      "                                                          \"['LAY_AND_WATCHED', \"\n",
      "                                                          \"'Birthday']\",\n",
      "                                                          \"['SHUFFLED_OFF_INTO', \"\n",
      "                                                          \"'Kitchen']\",\n",
      "                                                          \"['MOVED_OUT_OF', \"\n",
      "                                                          \"'Cupboard']\",\n",
      "                                                          \"['MOVED_INTO', \"\n",
      "                                                          '\"Dudley\\'s second '\n",
      "                                                          'bedroom\"]',\n",
      "                                                          \"['COULD_SEE', \"\n",
      "                                                          '\"Uncle vernon\\'s '\n",
      "                                                          'shiny black shoes\"]',\n",
      "                                                          \"['RECEIVED', 'First \"\n",
      "                                                          \"letter']\",\n",
      "                                                          \"['RECEIVED', \"\n",
      "                                                          \"'Letter']\",\n",
      "                                                          \"['PICKED_UP', \"\n",
      "                                                          \"'Letter']\",\n",
      "                                                          \"['STARED_AT', \"\n",
      "                                                          \"'Letter']\",\n",
      "                                                          \"['EARNED', \"\n",
      "                                                          \"'Longest-ever \"\n",
      "                                                          \"punishment']\",\n",
      "                                                          \"['DREAMED_OF', \"\n",
      "                                                          \"'Unknown relation']\",\n",
      "                                                          \"['FELL_HARD_ON', \"\n",
      "                                                          \"'Concrete floor']\",\n",
      "                                                          \"['IMAGINED', \"\n",
      "                                                          \"'School']\",\n",
      "                                                          \"['SITTING_IN', \"\n",
      "                                                          \"'Car']\",\n",
      "                                                          \"['HATED', 'Mrs. \"\n",
      "                                                          \"figg']\",\n",
      "                                                          \"['TURNING_OVER', \"\n",
      "                                                          \"'Bacon']\",\n",
      "                                                          \"['WORE', 'Glasses']\",\n",
      "                                                          \"['SLEPT_IN', \"\n",
      "                                                          \"'Cupboard']\",\n",
      "                                                          \"['HAD', 'Dream']\",\n",
      "                                                          \"['HAD', 'Sleepless \"\n",
      "                                                          \"night']\",\n",
      "                                                          \"['HAD', 'Scar']\"],\n",
      "                                          'kg_schema': {'schema': 'Node '\n",
      "                                                                  'properties '\n",
      "                                                                  'are the '\n",
      "                                                                  'following:\\n'\n",
      "                                                                  'Entity {id: '\n",
      "                                                                  'STRING}\\n'\n",
      "                                                                  'Relationship '\n",
      "                                                                  'properties '\n",
      "                                                                  'are the '\n",
      "                                                                  'following:\\n'\n",
      "                                                                  '\\n'\n",
      "                                                                  'The '\n",
      "                                                                  'relationships '\n",
      "                                                                  'are the '\n",
      "                                                                  'following:\\n'\n",
      "                                                                  \"(:Entity)-[:WAS_DIRECTOR_OF]->(:Entity),(:Entity)-[:HAD]->(:Entity),(:Entity)-[:CAME_INTO]->(:Entity),(:Entity)-[:FELL_ASLEEP_QUICKLY]->(:Entity),(:Entity)-[:SHUDDERED_TO_THINK]->(:Entity),(:Entity)-[:HATED]->(:Entity),(:Entity)-[:WERE]->(:Entity),(:Entity)-[:ARRIVED_IN]->(:Entity),(:Entity)-[:SAT_WITH]->(:Entity),(:Entity)-[:YELLED_AT]->(:Entity),(:Entity)-[:STOPPED]->(:Entity),(:Entity)-[:SAW]->(:Entity),(:Entity)-[:LAY_AWAKE]->(:Entity),(:Entity)-[:HAD_GOTTEN_INTO]->(:Entity),(:Entity)-[:FOUNDED_IN]->(:Entity),(:Entity)-[:IS]->(:Entity),(:Entity)-[:HAVE_BEEN_BEHAVING]->(:Entity),(:Entity)-[:HUNT_AT]->(:Entity),(:Entity)-[:KICKED]->(:Entity),(:Entity)-[:SLEPT_IN]->(:Entity),(:Entity)-[:WAS_HOWLING]->(:Entity),(:Entity)-[:LEARNED]->(:Entity),(:Entity)-[:UNWRAPPED]->(:Entity),(:Entity)-[:DANGLING_OVER]->(:Entity),(:Entity)-[:BAWLING_AT]->(:Entity),(:Entity)-[:SNORED]->(:Entity),(:Entity)-[:WAS]->(:Entity),(:Entity)-[:BROKE]->(:Entity),(:Entity)-[:THOUGHT]->(:Entity),(:Entity)-[:CELEBRATING]->(:Entity),(:Entity)-[:WAS_SHOWING]->(:Entity),(:Entity)-[:SLIPPED]->(:Entity),(:Entity)-[:TURNED]->(:Entity),(:Entity)-[:SAID]->(:Entity),(:Entity)-[:PATTED]->(:Entity),(:Entity)-[:TOLD]->(:Entity),(:Entity)-[:TOOK]->(:Entity),(:Entity)-[:LAID]->(:Entity),(:Entity)-[:WEARING]->(:Entity),(:Entity)-[:ASKED]->(:Entity),(:Entity)-[:BLEW]->(:Entity),(:Entity)-[:FLINCHED]->(:Entity),(:Entity)-[:PULLED_OUT]->(:Entity),(:Entity)-[:EYED]->(:Entity),(:Entity)-[:TURNED_UP_IN]->(:Entity),(:Entity)-[:HOLDING]->(:Entity),(:Entity)-[:SWUNG]->(:Entity),(:Entity)-[:SOBBED]->(:Entity),(:Entity)-[:MOVED_OUT_OF]->(:Entity),(:Entity)-[:SHUFFLED_OFF_INTO]->(:Entity),(:Entity)-[:LAY_AND_WATCHED]->(:Entity),(:Entity)-[:RAN]->(:Entity),(:Entity)-[:HEART_TWANGING_LIKE]->(:Entity),(:Entity)-[:ANGRILY]->(:Entity),(:Entity)-[:TURNED_OUT_TO_BE]->(:Entity),(:Entity)-[:WORE]->(:Entity),(:Entity)-[:TURNING_OVER]->(:Entity),(:Entity)-[:SITTING_IN]->(:Entity),(:Entity)-[:IMAGINED]->(:Entity),(:Entity)-[:FELL_HARD_ON]->(:Entity),(:Entity)-[:DREAMED_OF]->(:Entity),(:Entity)-[:EARNED]->(:Entity),(:Entity)-[:STARED_AT]->(:Entity),(:Entity)-[:PICKED_UP]->(:Entity),(:Entity)-[:RECEIVED]->(:Entity),(:Entity)-[:COULD_SEE]->(:Entity),(:Entity)-[:MOVED_INTO]->(:Entity),(:Entity)-[:WAS_ASLEEP_AT]->(:Entity),(:Entity)-[:WAS_NO_LONGER]->(:Entity),(:Entity)-[:ENTERED]->(:Entity),(:Entity)-[:HELD]->(:Entity),(:Entity)-[:HAD_BEEN_LYING_AT]->(:Entity),(:Entity)-[:SLAMMED]->(:Entity),(:Entity)-[:PARKED_AT]->(:Entity),(:Entity)-[:WAS_CARRYING]->(:Entity),(:Entity)-[:WAS_SAYING_IN]->(:Entity),(:Entity)-[:SNAPPED]->(:Entity),(:Entity)-[:READ]->(:Entity),(:Entity)-[:HAD_BEEN_CHASING]->(:Entity),(:Entity)-[:ATE_IN]->(:Entity),(:Entity)-[:DIDN'T_STOP_TO]->(:Entity),(:Entity)-[:JABBED_ITS_TAIL_AT]->(:Entity),(:Entity)-[:WAS_IN]->(:Entity),(:Entity)-[:DID_GO]->(:Entity),(:Entity)-[:CAME_THE_SOUND_OF]->(:Entity),(:Entity)-[:CAME]->(:Entity),(:Entity)-[:FOUND_THEIR_WAY_INTO]->(:Entity),(:Entity)-[:CAME_OVER_TO]->(:Entity),(:Entity)-[:BLEW_UP_AROUND]->(:Entity),(:Entity)-[:WAS_DANGLING_OVER]->(:Entity),(:Entity)-[:BECOMES_A_MEMBER_OF]->(:Entity),(:Entity)-[:KNOWN_FOR]->(:Entity),(:Entity)-[:LOYALTY_TO]->(:Entity),(:Entity)-[:BECOMES]->(:Entity),(:Entity)-[:IN]->(:Entity),(:Entity)-[:BY]->(:Entity),(:Entity)-[:GODFATHER_OF]->(:Entity),(:Entity)-[:FALSELY_ACCUSED_OF]->(:Entity),(:Entity)-[:BELIEVED_TO_HAVE_BEEN_KILLED_BY]->(:Entity),(:Entity)-[:RETURNS_TO]->(:Entity),(:Entity)-[:TEACHES]->(:Entity),(:Entity)-[:ASSIST]->(:Entity),(:Entity)-[:TODAY]->(:Entity),(:Entity)-[:MARRIED_TO]->(:Entity),(:Entity)-[:IS_MEMBER_OF]->(:Entity),(:Entity)-[:HAS]->(:Entity)\"}}}\n"
     ]
    }
   ],
   "source": [
    "import pprint\n",
    "\n",
    "pp = pprint.PrettyPrinter()\n",
    "pp.pprint(response.metadata)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
